package cn.ycc1.my.rag.controller;

import com.alibaba.dashscope.common.ResultCallback;
import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.embeddings.TextEmbeddingParam;
import com.alibaba.dashscope.embeddings.TextEmbeddingResult;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import lombok.extern.java.Log;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.IOException;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.Semaphore;

/**
 * 向量化
 * @author ycc
 * @date 2025/4/20
 * https://bailian.console.aliyun.com/?tab=doc#/doc/?type=model&url=https%3A%2F%2Fhelp.aliyun.com%2Fdocument_detail%2F2842587.html
 */
@RestController
@RequestMapping("/embedding")
@Log
public class EmbeddingController {
    private static final int DASHSCOPE_MAX_BATCH_SIZE = 25;

    @Autowired
    ResourceLoader resourceLoader;

    /**
     * 输入字符串
     * @throws ApiException
     * @throws NoApiKeyException
     */
    @RequestMapping("/basicCall")
    public void basicCall() throws ApiException, NoApiKeyException {
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                .texts(Arrays.asList("风急天高猿啸哀", "渚清沙白鸟飞回", "无边落木萧萧下", "不尽长江滚滚来")).build();
        TextEmbedding textEmbedding = new TextEmbedding();
        TextEmbeddingResult result = textEmbedding.call(param);
        System.out.println(result);
    }

    /**
     * 异步调用
     * @throws ApiException
     * @throws NoApiKeyException
     * @throws InterruptedException
     */
    @RequestMapping("/callWithCallback")
    public static void callWithCallback() throws ApiException, NoApiKeyException, InterruptedException{
        TextEmbeddingParam param = TextEmbeddingParam
                .builder()
                .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                .texts(Arrays.asList("风急天高猿啸哀", "渚清沙白鸟飞回", "无边落木萧萧下", "不尽长江滚滚来")).build();
        TextEmbedding textEmbedding = new TextEmbedding();
        Semaphore sem = new Semaphore(0);
        textEmbedding.call(param, new ResultCallback<TextEmbeddingResult>() {

            @Override
            public void onEvent(TextEmbeddingResult message) {
                System.out.println(message);
            }
            @Override
            public void onComplete(){
                sem.release();
            }

            @Override
            public void onError(Exception err){
                System.out.println(err.getMessage());
                err.printStackTrace();
                sem.release();
            }

        });
        sem.acquire();
    }

    /**
     * 输入字符串列表，分批次调用
     */
    @RequestMapping("/callList")
    public void callList() {
        List<String> inputs = Arrays.asList(
                "风急天高猿啸哀",
                "渚清沙白鸟飞回",
                "无边落木萧萧下",
                "不尽长江滚滚来"
        );

        TextEmbeddingResult result = null;
        int batchCounter = 0;

        for (int i = 0; i < inputs.size(); i += DASHSCOPE_MAX_BATCH_SIZE) {
            List<String> batch = inputs.subList(i, Math.min(i + DASHSCOPE_MAX_BATCH_SIZE, inputs.size()));
            TextEmbeddingParam param = TextEmbeddingParam.builder()
                    .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                    .texts(batch)
                    .build();

            TextEmbedding textEmbedding = new TextEmbedding();
            try {
                TextEmbeddingResult resp = textEmbedding.call(param);
                if (resp != null) {
                    if (result == null) {
                        result = resp;
                    } else {
                        for (var emb : resp.getOutput().getEmbeddings()) {
                            emb.setTextIndex(emb.getTextIndex() + batchCounter);
                            result.getOutput().getEmbeddings().add(emb);
                        }
                        result.getUsage().setTotalTokens(result.getUsage().getTotalTokens() + resp.getUsage().getTotalTokens());
                    }
                } else {
                    System.out.println(resp);
                }
            } catch (ApiException | NoApiKeyException e) {
                e.printStackTrace();
            }
            batchCounter += batch.size();
        }

        System.out.println(result);
//        log.info("Embeddings: {}", result.toString());
    }

    /**
     * 输入文件
     */
    @RequestMapping("/callFile")
    public void callFile() throws IOException {

        Resource resource = resourceLoader.getResource("classpath:texts_to_embedding.txt");
        if (resource == null) {
            System.out.println("文件不存在");
            return;
        }
        String filePath = resource.getFile().getAbsolutePath();
        System.out.println(filePath);
//        try (BufferedReader reader = new BufferedReader(new FileReader("<文件所来自的内容根的路径>"))) {
        try (BufferedReader reader = new BufferedReader(new FileReader(filePath))) {
            StringBuilder content = new StringBuilder();
            String line;
            while ((line = reader.readLine()) != null) {
                content.append(line).append("\n");
            }

            TextEmbeddingParam param = TextEmbeddingParam.builder()
                    .model(TextEmbedding.Models.TEXT_EMBEDDING_V3)
                    .text(content.toString())
                    .build();

            TextEmbedding textEmbedding = new TextEmbedding();
            TextEmbeddingResult result = textEmbedding.call(param);

            if (result != null) {
                System.out.println(result);
            } else {
                System.out.println("Failed to get embedding: " + result);
            }
        } catch (IOException | ApiException | NoApiKeyException e) {
            e.printStackTrace();
        }
    }
}
